import base64
from io import BytesIO

import cv2
import numpy as np

from math import *
from pathlib import Path
from typing import Union, Tuple, Callable, Literal, Sequence
from PIL import Image
from matplotlib import pyplot as plt

from .decorator import param2path
from .log import logger


class ImageTools(object):
    """
    图片处理类
    """

    @param2path({'img_path': 1})
    def __init__(self, img_path: Union[str, Path]):
        """
        初始化
        """
        self._path_c = img_path
        self._path = str(self._path_c)
        # 获取图片数据
        logger.info(f"图片地址：{self._path_c}")
        self._img = cv2.imread(self._path)
        if self._img is None:
            _image = Image.open(self._path)
            self._img = cv2.cvtColor(np.asarray(_image), cv2.COLOR_RGB2BGR)
        logger.info(f"原始尺寸(h, w, c)：{self.img.shape}")
        self.center_offset = (0, 0)  # 中心偏移量
        self.center_rotate = self.center  # 旋转中心
        self.fill_num = (0,0,0,0)  # 填充的数值

    @property
    def base_path_c(self) -> Path:
        """
        图片的文件夹
        Returns:

        """
        return Path(self.path_c.parent)

    @property
    def center(self) -> Tuple:
        """
        图片中心
        Returns:

        """
        _x, _y = self.size
        return _x / 2, _y / 2

    @property
    def size(self) -> Tuple:
        """
        图片尺寸 (w, h)
        Returns:

        """
        return self.img.shape[:2][::-1]

    @property
    def path_c(self) -> Path:
        """
        图片源地址对象
        Returns:

        """
        return self._path_c

    @property
    def img(self) -> np.ndarray:
        """
        图片数据
        Returns:

        """
        return self._img

    @property
    def img_rgb(self) -> np.ndarray:
        """
        图片数据
        Returns:

        """
        return cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)

    @property
    def to_base64(self) -> str:
        pil_img = Image.fromarray(self.img_rgb)
        buff = BytesIO()
        pil_img.save(buff, format="JPEG")
        return base64.b64encode(buff.getvalue()).decode('utf-8')

    def execute(self, call_back: Callable):
        """
        执行
        Args:
            call_back: 回调

        Returns:

        """
        if ret := call_back(self.img):
            self._img = ret
        return self

    def crop(self, start: Tuple = (0, 0), end: Tuple = None) -> 'ImageTools':
        """
        图片裁剪
        Args:
            start: 开始坐标
            end: 结束坐标，默认值为图片尺寸

        Returns:

        """
        _max_size = self.img.shape[:2][::-1]
        _x, _y = start
        # 结尾坐标
        end = list(end)
        if not end:
            end = _max_size
        for _i, (_max, _d) in enumerate(zip(_max_size, end)):
            if not _d:
                end[_i] = _max
            elif _d < 0:
                end[_i] = _max + _d
        _w, _h = end

        self._img = self.img[_y: _y + _h, _x: _x + _w]
        logger.info(f"裁剪后尺寸(h, w, c)：{self.img.shape}")
        return self

    def resize(self, d_size: Tuple = None, max_side: int = 0, min_side: int = 0,
               ratio: Union[int, float, Tuple] = 1) -> 'ImageTools':
        """
        改变图片大小
        Args:
            d_size: 指定修改尺寸 (w, h)
            max_side: 最大边长 当 d_size 不合法时启用
            min_side: 最小边长 当 d_size 不合法时启用
            ratio: 比例 当 d_size 不合法时启用

        Returns:

        """
        img_w, img_h = self.size  # 获取图片宽高
        if img_w > img_h:
            max_index = 0
            min_index = 1
        else:
            max_index = 1
            min_index = 0

        # 指定最大最小边长
        if not d_size:
            d_size = [None, None]
            if max_side > 0:
                d_size[max_index] = max_side
            if min_side > 0:
                d_size[min_index] = min_side

        if d_size:
            _w, _h = d_size
            if not _w and _h:
                d_size = (int(_h * img_w / img_h), _h)
            if not _h and _w:
                d_size = (_w, int(_w * img_h / img_w))
            self._img = cv2.resize(self.img, tuple(d_size))
        else:
            if isinstance(ratio, tuple):
                _fx, _fy = ratio
            else:
                _fx, _fy = [ratio] * 2
            self._img = cv2.resize(self.img, (0, 0), fx=_fx, fy=_fy)
        logger.info(f"修改后尺寸(h, w, c)：{self.img.shape}")
        return self

    def rotate(self, angle: float, scale: float = 1,
               point: Tuple = None) -> 'ImageTools':
        """
        旋转
        Args:
            angle: 角度
            scale: 缩放比例
            point: 旋转中心，默认为图片中心， 暂时不支持非中心旋转

        Returns:

        """
        _w, _h = self.size  # 图片尺寸
        _h_n = int(_w * fabs(sin(radians(angle))) + _h * fabs(
            cos(radians(angle))))
        _w_n = int(_h * fabs(sin(radians(angle))) + _w * fabs(
            cos(radians(angle))))
        if not point:
            point = self.center_rotate
        else:
            self.center_rotate = point
        # 第一个参数旋转中心，第二个参数旋转角度，第三个参数：缩放比例
        m = cv2.getRotationMatrix2D(point, angle, scale)
        # 中心偏移
        self.center_offset = ((_w_n - _w) / 2, (_h_n - _h) / 2)
        m[0, 2] += self.center_offset[0]
        m[1, 2] += self.center_offset[1]

        self._img = cv2.warpAffine(
            self.img, m,
            (_w_n, _h_n),
            borderValue=(255, 255, 255))
        logger.info(f"角度：{angle}，缩放：{scale}，中心：{point}")
        logger.info(f"旋转后尺寸(h, w, c)：{self.img.shape}")
        return self

    def draw_point(self, points: Sequence[Tuple[int, int]]
                   ) -> 'ImageTools':
        """
        在图片上描点
        Args:
            points: 坐标的序列，只支持整数坐标

        Returns:

        """
        _r = max(self.size) // 100  # 点的半径
        for _x, _y in points:
            self._img = cv2.circle(
                self.img, (int(_x), int(_y)), _r,
                (255, 0, 0), thickness=-1, lineType=cv2.FILLED)
        return self

    def show(self, t: Literal['cv', 'plt', 'pil'] = 'pil') -> 'ImageTools':
        """
        图片展示
        Args:
            t: 展示形式
        Returns:

        """
        if t == 'cv':
            cv2.imshow(self.path_c.name, self.img)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        elif t == 'pil':
            _img = Image.fromarray(self.img_rgb)
            _img.show()
        elif t == 'plt':
            plt.imshow(self.img)
            plt.show()
        else:
            pass
        return self

    @param2path({'to_path_c': 1}, ['to_path_c'])
    def save(self, to_path_c: Union[Path, str] = r'.',
             stem: str = '') -> 'ImageTools':
        """
        保存
        Args:
            to_path_c: 保存到的文件夹地址
            stem: 文件名（不包含后缀）
        Returns:

        """
        to_file_path = str(
            to_path_c.joinpath(
                f"{stem if stem else self.path_c.stem.lower()}"
                f"{self.path_c.suffix.lower()}")
        )

        logger.info(f"保存地址：{to_file_path}")
        cv2.imwrite(to_file_path, self.img)
        return self

    def fill(self, size: Union[Tuple[float, float], float], is_avg: bool = True,
             **kwargs) -> 'ImageTools':
        """
        填充
        Args:
            size: 填充后图片尺寸(w, h)
            is_avg: 是否平均填充
            **kwargs: 填充参数

        Returns:

        """

        def _fill_size(t: float, s: float) -> Tuple:
            _pre = 0
            _suf = 0
            if t > s:
                logger.info(f"{'平均' if is_avg else '尾部'}填充：{s} -> {t}")
                _diff = t - s
                if is_avg:
                    _pre = floor(_diff / 2)
                    _suf = ceil(_diff / 2)
                else:
                    _suf = _diff
            return _pre, _suf

        if str(size).isdigit():
            size = (size, size)
        _w_t, _h_t = size  # 目标宽高
        _w, _h = self.size  # 原始宽高
        _l, _r = _fill_size(_w_t, _w)
        _t, _b = _fill_size(_h_t, _h)
        self.fill_num = (_t, _b, _l, _r)
        _param = {
            'borderType': cv2.BORDER_CONSTANT,
            'value': (192, 192, 192),
        }
        _param.update(kwargs)
        # 填充
        self._img = cv2.copyMakeBorder(self._img, _t, _b, _l, _r, **_param)
        return self
